import torch
import numpy as np
from torch.utils.data import Dataset

class HiCLoopDataset(Dataset):
    """Hi-C Loop检测数据集"""
    def __init__(self, hic_data, genomic_features, labels, metadata, 
                 resolution=5000, transform=None):
        """
        hic_data: Hi-C交互矩阵字典 {chrom: np.array(seq_len, seq_len)}
        genomic_features: 多组学特征字典 {chrom: {
            'positions': np.array(seq_len, 3),  # [bin_start, bin_end, bin_center]
            'ctcf': np.array(seq_len),         # CTCF信号强度
            'h3k27ac': np.array(seq_len),      # H3K27ac信号强度
            'h3k4me3': np.array(seq_len),      # H3K4me3信号强度
            'promoter': np.array(seq_len),     # 启动子密度
            'enhancer': np.array(seq_len),     # 增强子密度
            'expression': np.array(seq_len)    # 基因表达水平
        }}
        labels: 标签字典 {chrom: np.array(seq_len, seq_len)}  # Loop位置二元矩阵
        metadata: 元数据字典 {
            'tad_boundaries': {chrom: [(start, end), ...]},  # TAD边界
            'promoter_regions': {chrom: [(start, end), ...]}, # 启动子区域
            'chrom_sizes': {chrom: size}  # 染色体大小
        }
        resolution: Hi-C分辨率(bp)
        transform: 数据增强函数
        """
        self.hic_data = hic_data
        self.genomic_features = genomic_features
        self.labels = labels
        self.metadata = metadata
        self.resolution = resolution
        self.transform = transform
        
        # 预计算所有样本索引
        self.samples = []
        for chrom in hic_data.keys():
            seq_len = hic_data[chrom].shape[0]
            # 创建滑动窗口样本 (步长为窗口大小的1/2)
            window_size = 256  # 每个样本包含256个bin
            step = window_size // 2
            for start_idx in range(0, seq_len - window_size + 1, step):
                end_idx = start_idx + window_size
                self.samples.append({
                    'chrom': chrom,
                    'start_idx': start_idx,
                    'end_idx': end_idx
                })
    
    def __len__(self):
        return len(self.samples)
    
    def __getitem__(self, idx):
        sample_info = self.samples[idx]
        chrom = sample_info['chrom']
        start_idx = sample_info['start_idx']
        end_idx = sample_info['end_idx']
        
        # 1. 提取Hi-C子矩阵
        hic_sub = self.hic_data[chrom][start_idx:end_idx, start_idx:end_idx]
        
        # 2. 提取基因组特征
        genomic_sub = {
            k: v[start_idx:end_idx] for k, v in self.genomic_features[chrom].items()
            if k != 'positions'  # 位置信息单独处理
        }
        
        # 3. 提取位置信息
        positions = self.genomic_features[chrom]['positions'][start_idx:end_idx]
        
        # 4. 提取标签
        label_sub = self.labels[chrom][start_idx:end_idx, start_idx:end_idx]
        
        # 5. 创建位置差异矩阵
        bin_indices = np.arange(start_idx, end_idx)
        position_diff = np.abs(bin_indices[:, None] - bin_indices[None, :])
        
        # 6. 创建显著性掩码 (使用泊松分布)
        significance_mask = self._calculate_significance(hic_sub)
        
        # 7. 创建功能区域特征
        functional_features = self._create_functional_features(
            chrom, start_idx, end_idx
        )
        
        # 8. 数据增强
        if self.transform:
            hic_sub, genomic_sub, label_sub, significance_mask = self.transform(
                hic_sub, genomic_sub, label_sub, significance_mask
            )
        
        # 转换为张量
        hic_tensor = torch.tensor(hic_sub, dtype=torch.float32)
        genomic_tensor = torch.stack([
            torch.tensor(genomic_sub['ctcf'], dtype=torch.float32),
            torch.tensor(genomic_sub['h3k27ac'], dtype=torch.float32),
            torch.tensor(genomic_sub['h3k4me3'], dtype=torch.float32),
            torch.tensor(genomic_sub['promoter'], dtype=torch.float32),
            torch.tensor(genomic_sub['enhancer'], dtype=torch.float32),
            torch.tensor(genomic_sub['expression'], dtype=torch.float32)
        ], dim=1)  # [seq_len, 6]
        
        label_tensor = torch.tensor(label_sub, dtype=torch.float32)
        position_diff_tensor = torch.tensor(position_diff, dtype=torch.int32)
        significance_mask_tensor = torch.tensor(significance_mask, dtype=torch.float32)
        functional_tensor = torch.tensor(functional_features, dtype=torch.float32)
        
        return {
            'hic_matrix': hic_tensor,
            'genomic_features': genomic_tensor,
            'label': label_tensor,
            'position_diff': position_diff_tensor,
            'significance_mask': significance_mask_tensor,
            'functional_features': functional_tensor,
            'chrom': chrom,
            'start_idx': start_idx,
            'end_idx': end_idx
        }
    
    def _calculate_significance(self, hic_matrix):
        """使用泊松分布计算显著性掩码"""
        # 简化的泊松显著性检测 (实际实现需要完整逻辑)
        local_radius = 30
        seq_len = hic_matrix.shape[0]
        significance_mask = np.zeros((seq_len, seq_len))
        
        for i in range(seq_len):
            for j in range(max(0, i-local_radius), min(seq_len, i+local_radius+1)):
                # 获取局部区域
                i_start = max(0, i - local_radius)
                i_end = min(seq_len, i + local_radius + 1)
                j_start = max(0, j - local_radius)
                j_end = min(seq_len, j + local_radius + 1)
                
                local_region = hic_matrix[i_start:i_end, j_start:j_end].copy()
                # 排除中心点
                local_region[i-i_start, j-j_start] = 0
                
                # 计算局部背景
                local_mean = np.mean(local_region)
                # 简化的显著性判断
                if hic_matrix[i, j] > 3 * local_mean:
                    significance_mask[i, j] = 1
                    significance_mask[j, i] = 1
        
        return significance_mask
    
    def _create_functional_features(self, chrom, start_idx, end_idx):
        """创建功能区域特征"""
        seq_len = end_idx - start_idx
        functional_features = np.zeros((seq_len, 3))  # TAD边界, 启动子, 增强子
        
        # 1. TAD边界特征
        tad_boundaries = self.metadata['tad_boundaries'].get(chrom, [])
        for start, end in tad_boundaries:
            # 转换为bin索引
            start_bin = start // self.resolution
            end_bin = end // self.resolution
            
            # 标记边界区域
            if start_idx <= start_bin < end_idx:
                idx = start_bin - start_idx
                functional_features[idx, 0] = 1.0
            if start_idx <= end_bin < end_idx:
                idx = end_bin - start_idx
                functional_features[idx, 0] = 1.0
        
        # 2. 启动子特征
        promoter_regions = self.metadata['promoter_regions'].get(chrom, [])
        for start, end in promoter_regions:
            start_bin = start // self.resolution
            end_bin = end // self.resolution
            
            # 标记启动子区域内的所有bin
            for bin_idx in range(max(start_idx, start_bin), min(end_idx, end_bin+1)):
                idx = bin_idx - start_idx
                functional_features[idx, 1] = 1.0
        
        # 3. 增强子特征 (从基因组特征中提取)
        # 实际应用中可能已有增强子标注
        
        return functional_features

# 数据增强函数示例
class RandomStrideAugment:
    """随机步长数据增强 - 模拟不同分辨率"""
    def __init__(self, stride_range=(1, 3)):
        self.stride_range = stride_range
        
    def __call__(self, hic, genomic, label, mask):
        stride = np.random.randint(*self.stride_range)
        
        # 降采样
        hic_sub = hic[::stride, ::stride]
        genomic_sub = {k: v[::stride] for k, v in genomic.items()}
        label_sub = label[::stride, ::stride]
        mask_sub = mask[::stride, ::stride]
        
        # 上采样回原始尺寸
        def upsample(array, stride):
            return np.repeat(np.repeat(array, stride, axis=0), stride, axis=1)
        
        hic_sub = upsample(hic_sub, stride)[:hic.shape[0], :hic.shape[1]]
        label_sub = upsample(label_sub, stride)[:label.shape[0], :label.shape[1]]
        mask_sub = upsample(mask_sub, stride)[:mask.shape[0], :mask.shape[1]]
        
        genomic_sub = {k: np.repeat(v, stride)[:len(v)] for k, v in genomic_sub.items()}
        
        return hic_sub, genomic_sub, label_sub, mask_sub